from __main__ import *
from io import TextIOWrapper
import os
import logging
from Disassembly import DisassemblyDecoder
from TraceDecode import TraceDecode, TraceCmd, ParserData


class TraceDumpDecoder:
    def __init__(self, disassemble_file: str, trace_decode_file: str, output_file: str):
        self.__disassemble = DisassemblyDecoder(disassemble_file)
        self.__trace_decode = TraceDecode(trace_decode_file)
        self.__output = open(output_file, "w")

    def __dumpKeyword(self, trace: ParserData):
        return trace.cmd + "\n"

    def __dumpInstr(
        self,
        trace: ParserData,
        instr: list[list[str]],
    ):
        output = ""
        cmd = TraceCmd(trace.cmd)
        if cmd.exec:
            actual_instruction = instr[-1][1].strip()
            if actual_instruction == cmd:
                for i in instr:
                    output = output + "0x" + i[0].strip() + "\t" + i[1].strip() + "\n"
            else:
                logging.error(
                    "Unexpected end instruction "
                    + str(trace.cmd)
                    + " instructions from address "
                    + trace.addr
                    + ". Please check ELF file provided is for the snapshots and binaries provided."
                )
                return None
        else:
            for i in instr[:-1]:
                output = output + "0x" + i[0].strip() + "\t" + i[1].strip() + "\n"
        return output

    def __dumpDecodeHandle(self, file: TextIOWrapper):
        trace = self.__trace_decode.getTrace()
        output = ""
        for t in trace:
            disa_instr = self.__disassemble.get(t.saddr, t.eaddr)
            if disa_instr:
                out = self.__dumpInstr(t, disa_instr)
                if out is None:
                    break
                output += out
            else:
                output += self.__dumpKeyword(t)
        file.write(output)
        file.close()

    def dumpDecode(self):
        self.__dumpDecodeHandle(self.__output)
